package ace.cmp.spring.data.jdbc.querydsl.annotation.processor;

import com.querydsl.apt.Configuration;
import com.querydsl.apt.ExtendedTypeFactory;
import com.querydsl.apt.TypeElementHandler;
import com.querydsl.codegen.EntityType;
import com.querydsl.codegen.QueryTypeFactory;
import com.querydsl.codegen.TypeMappings;
import com.querydsl.sql.codegen.support.PrimaryKeyData;
import java.util.Arrays;
import java.util.List;
import java.util.stream.Collectors;
import javax.lang.model.element.Element;
import javax.lang.model.element.TypeElement;
import org.springframework.data.annotation.Id;

/**
 * @author caspar
 * @date 2023/4/5 10:39
 */
public class TypeWithPrimaryKeyElementHandler extends TypeElementHandler {
  public static final String PRIMARY_KEY_VAR_NAME = "primaryKey";

  public TypeWithPrimaryKeyElementHandler(
      Configuration configuration,
      ExtendedTypeFactory typeFactory,
      TypeMappings typeMappings,
      QueryTypeFactory queryTypeFactory) {
    super(configuration, typeFactory, typeMappings, queryTypeFactory);
  }

  @Override
  public EntityType handleEntityType(TypeElement element) {
    EntityType entityType = super.handleEntityType(element);

    List<? extends Element> elements = element.getEnclosedElements();

    List<String> primaryKeyColumns =
        elements.stream()
            .map(
                childElement -> {
                  Id idAnnotation = childElement.getAnnotation(Id.class);

                  if (idAnnotation != null) {

                    String name = childElement.getSimpleName().toString();
                    String column = name;
                    return column;
                  }
                  return null;
                })
            .filter(p -> p != null)
            .collect(Collectors.toList());
    if (primaryKeyColumns.size() > 0) {
      PrimaryKeyData primaryKeyData =
          new PrimaryKeyData(PRIMARY_KEY_VAR_NAME, primaryKeyColumns.toArray(new String[0]));
      entityType.getData().put(PrimaryKeyData.class, Arrays.asList(primaryKeyData));
    }

    return entityType;
  }
}
